import datetime
import logging
import os

from paddleocr import PaddleOCR
import re
import time
from werkzeug.utils import secure_filename

import torch
from torchvision import transforms
from PIL import Image

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 加密
aes_key = b'L03-Cumminschine'
iv_key = b'Cty-china-ER-app'

# 图片路径
image_file_path = "./images/"
model_path = r"./models/model_draft_1129.pth"
loaded_model = torch.load(model_path, map_location=torch.device('cpu'))
ocr = PaddleOCR(
    use_angle_cls=True,
    cls_model_dir=r"./models/cls2",
    det_model_dir=r"./models/det2",
    rec_model_dir=r"./models/rec2"
)
# 角度不对 模糊 假的  未识别
class_labels = ['您拍摄照片的角度倾斜', '您拍摄的照片比较模糊', '您拍摄的照片可能存在翻拍', '']


# 上传文件
# type (1  ESN) (2 OEMName、ChassisNumber、VIN) (3 OEMName、VIN) (4、WarrantyStartDate、ChassisNumber) (5、WarrantyStartDate)
def upload(file_content, file_name, type):
    start_time = time.time()
    logger.info("begin: {:.2f}s".format(start_time))

    # type = int(headers.get("type"))
    # filename = secure_filename(upload_file.filename)
    file_path = image_file_path + file_name
    logger.info(f"saving file into : {file_path}")
    __save_file(file_path, file_content)

    save_time = time.time()
    ocr_result = __run_model(file_path, type)
    msg = ''
    if ocr_result is None:
        msg = __single_cf_test(file_path)

    logger.info("识别内容" + str(ocr_result))
    ocr_time = time.time()
    time_ocr = "ocrTime: " + str(ocr_time - save_time)
    logger.info(str(time_ocr))
    print("allTime: {:.2f}s".format(ocr_time - start_time))
    # print(ocr)
    return {"code": 1, "data": ocr_result, 'msg': msg}


def __run_model(img_path, type):

    result = ocr.ocr(img_path, cls=True)
    logger.info("解析前：" + str(result))
    if result and len(result) > 0:
        result_arr = [line[1][0] for line in result[0]]
        logger.info("图片list：" + str(result_arr))
        return __resolve_ocr_result(result_arr, type)
    else:
        logger.warning("解析失败")
        return None


def __resolve_ocr_result(result_arr, type):
    ocr_results = {}
    # logger.info(f"type: {type}")
    # validate model outputs
    for i in range(len(result_arr)):
        v_vin = re.findall("[L]{1}[A-Z0-9]{16}", str(result_arr[i]))
        v_esn = re.findall("[7]{1}[0-9]{7}", str(result_arr[i]))
        v_chassis = re.findall("\d{4}-\d{2}-\d{2}", str(result_arr[i]))
        v_name = re.findall("\d{4}-\d{2}-\d{2}", str(result_arr[i]))
        v_date = re.findall("\d{4}-\d{2}-\d{2}", str(result_arr[i]))
        logger.info(f"v_vin: {v_vin}")
        if v_vin and type in [2, 3]:
            ocr_results["VIN"] = v_vin[0]

        if v_name and type in [2, 3]:
            ocr_results["OEMName"] = v_name[0]

        if v_esn and type == 1:
            ocr_results["ESN"] = v_esn[0]

        if v_chassis and type in [2, 4]:
            ocr_results["ChassisNumber"] = v_chassis[0]

        if v_date and type in [4, 5]:
            ocr_results["WarrantyStartDate"] = v_date[0]

    return ocr_results
# 识别失败-报错处理
def __single_cf_test(image_path):
    image = Image.open(image_path).convert('RGB')
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    input_tensor = transform(image)
    input_batch = input_tensor.unsqueeze(0)
    loaded_model.eval()
    with torch.no_grad():
        output = loaded_model(input_batch)

    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    predicted_class = torch.argmax(probabilities).item()
    predicted_label = class_labels[predicted_class]
    logger.info(f'--- \n Predicted Class: {predicted_label}')
    print(f'Class Probability: {probabilities[predicted_class].item()}')
    return predicted_label
    # classification result========================


# save file to local path
def __save_file(file_path, file_content):
    try:
        with open(file_path, "wb") as lf:
            lf.write(file_content)
        return True
    except Exception as e:
        logger.error('成功保存失败发生错误：', str(e))
        return False



if __name__ == '__main__':
    event = {
        "headers": {
            "type": "2"
        },
        # base64 encoded image
        "body": "123"
    }

    headers = event.get("headers")
    body = event.get('body')
    type = int(headers.get("type"))
    if type is None:
        raise ValueError("type is invalid")
    if body is None:
        raise ValueError("body is invalid")

    with open("./test.png", "rb") as img:
        file_content = img.read()

    result = upload(file_content, "test.png", type)
    print(f"result: {result}")
